#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <fcntl.h>
#include <sys/epoll.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <pthread.h>
#include <errno.h>

#define MAX_EVENTS 64
#define SERVER_PORT 8080
#define WORKER_THREADS 4
#define BUFFER_SIZE 1024

// 线程参数结构
typedef struct {
    int epoll_fd;
    int thread_id;
} worker_args;

// 全局 epoll 实例 (主Reactor)
int main_epoll_fd;

// 设置非阻塞
void set_nonblocking(int fd) {
    int flags = fcntl(fd, F_GETFL, 0);
    fcntl(fd, F_SETFL, flags | O_NONBLOCK);
}

// 处理接收的数据
void handle_data(int fd) {
    char buffer[BUFFER_SIZE];
    ssize_t bytes_read;
    
    while ((bytes_read = read(fd, buffer, BUFFER_SIZE)) > 0) {
        printf("Received %zd bytes\n", bytes_read);
        write(fd, buffer, bytes_read);  // 回显数据
    }
    
    if (bytes_read == 0 || (bytes_read == -1 && errno != EAGAIN)) {
        close(fd);
        printf("Connection closed\n");
    }
}

// 工作线程函数 (从Reactor)
void *worker_thread(void *arg) {
    worker_args *args = (worker_args *)arg;
    struct epoll_event events[MAX_EVENTS];
    
    while (1) {
        int n = epoll_wait(args->epoll_fd, events, MAX_EVENTS, -1);
        for (int i = 0; i < n; i++) {
            if (events[i].events & EPOLLIN) {
                handle_data(events[i].data.fd);
            }
        }
    }
    
    return NULL;
}

// 主线程函数 (主Reactor)
int main() {
    struct epoll_event event;
    struct sockaddr_in addr;
    int server_fd, client_fd;
    
    // 创建 TCP socket
    if ((server_fd = socket(AF_INET, SOCK_STREAM, 0)) < 0) {
        perror("socket failed");
        exit(EXIT_FAILURE);
    }
    
    // 设置 SO_REUSEADDR
    int opt = 1;
    setsockopt(server_fd, SOL_SOCKET, SO_REUSEADDR, &opt, sizeof(opt));
    
    // 绑定端口
    memset(&addr, 0, sizeof(addr));
    addr.sin_family = AF_INET;
    addr.sin_addr.s_addr = INADDR_ANY;
    addr.sin_port = htons(SERVER_PORT);
    
    if (bind(server_fd, (struct sockaddr*)&addr, sizeof(addr)) < 0) {
        perror("bind failed");
        close(server_fd);
        exit(EXIT_FAILURE);
    }
    
    // 开始监听
    if (listen(server_fd, SOMAXCONN) < 0) {
        perror("listen failed");
        close(server_fd);
        exit(EXIT_FAILURE);
    }
    
    printf("Server listening on port %d\n", SERVER_PORT);
    
    // 创建主 epoll 实例
    if ((main_epoll_fd = epoll_create1(0)) < 0) {
        perror("epoll_create failed");
        close(server_fd);
        exit(EXIT_FAILURE);
    }
    
    // 添加服务器 socket 到主 epoll
    event.events = EPOLLIN;
    event.data.fd = server_fd;
    if (epoll_ctl(main_epoll_fd, EPOLL_CTL_ADD, server_fd, &event) < 0) {
        perror("epoll_ctl failed");
        close(server_fd);
        close(main_epoll_fd);
        exit(EXIT_FAILURE);
    }
    
    // 创建工作线程 (从Reactor)
    pthread_t threads[WORKER_THREADS];
    worker_args thread_args[WORKER_THREADS];
    int worker_epoll_fd[WORKER_THREADS];
    
    for (int i = 0; i < WORKER_THREADS; i++) 
	{
        if ((worker_epoll_fd[i] = epoll_create1(0)) < 0) 
		{
            perror("worker epoll_create failed");
            exit(EXIT_FAILURE);
        }
        
        thread_args[i].epoll_fd = worker_epoll_fd[i];
        thread_args[i].thread_id = i;
        
        if (pthread_create(&threads[i], NULL, worker_thread, &thread_args[i]) != 0) 
		{
            perror("pthread_create failed");
            exit(EXIT_FAILURE);
        }
        printf("Worker thread %d created\n", i);
    }
    
    // 主事件循环
    struct epoll_event events[MAX_EVENTS];
    int current_worker = 0;
    
    while (1) {
        int n = epoll_wait(main_epoll_fd, events, MAX_EVENTS, -1);
        
        for (int i = 0; i < n; i++) {
            // 处理新连接
            if (events[i].data.fd == server_fd) {
                struct sockaddr_in client_addr;
                socklen_t client_len = sizeof(client_addr);
                
                if ((client_fd = accept(server_fd, (struct sockaddr*)&client_addr, &client_len)) < 0) {
                    perror("accept failed");
                    continue;
                }
                
                set_nonblocking(client_fd);
                printf("New connection accepted\n");
                
                // 将新连接分配给工作线程 (轮询分配)
                event.events = EPOLLIN | EPOLLET;
                event.data.fd = client_fd;
                
                if (epoll_ctl(worker_epoll_fd[current_worker], EPOLL_CTL_ADD, client_fd, &event) < 0) {
                    perror("worker epoll_ctl failed");
                    close(client_fd);
                } else {
                    printf("Connection assigned to worker %d\n", current_worker);
                }
                
                current_worker = (current_worker + 1) % WORKER_THREADS;
            }
        }
    }
    
    // 清理 (实际上永远不会执行到这里)
    close(server_fd);
    close(main_epoll_fd);
    for (int i = 0; i < WORKER_THREADS; i++) {
        close(worker_epoll_fd[i]);
    }
    
    return 0;
}